package com.fc.connector.source;


import org.apache.flink.api.connector.source.*;
import org.apache.flink.core.io.SimpleVersionedSerializer;
import javax.annotation.Nullable;
import java.io.IOException;
import java.util.*;

//

/**
 * SplitEnumerator 实现（兼容 Flink 1.20）
 * 数据源分片枚举器
 *
 */
public class MySplitEnumerator implements SplitEnumerator<MySplit, MySplitEnumerator.CheckpointState> {
    private final SplitEnumeratorContext<MySplit> context;
    private final List<MySplit> pendingSplits = new ArrayList<>();
    private final Map<String, List<MySplit>> assignedSplits = new HashMap<>();
    private static final long DISCOVER_INTERVAL = 10000L; // 10秒分片发现间隔

    public MySplitEnumerator(SplitEnumeratorContext<MySplit> context, CheckpointState checkpoint) {
        this.context = context;
        if (checkpoint != null) {
            this.pendingSplits.addAll(checkpoint.getRemainingSplits());
        }
    }


    /**
     * 启动枚举器，这个方法调用一次，枚举器启动时调用。
     */

    @Override
    public void start() {
        // 异步分片发现（定期执行）
        context.callAsync(
                this::discoverNewSplits,
                (newSplits, error) -> {
                    if (error != null) return;
                    synchronized (this) {
                        //1、发现分片添加到待分配列表pendingSplits
                        pendingSplits.addAll(newSplits);
                        //2、分配分片
                        assignSplits();
                    }
                },
                0L, DISCOVER_INTERVAL
        );
    }



    private List<MySplit> discoverNewSplits() {
        // 实际应用中从外部系统（如 Kafka、文件系统）获取新分片
        return Collections.singletonList(new MySplit("split-" + UUID.randomUUID()));
    }

    // 分片分配逻辑（基于并行度哈希分配）:cite[1]:cite[6]
    private void assignSplits() {
        System.out.println("分片分配开始");
        if (pendingSplits.isEmpty()) return;

        Map<Integer, List<MySplit>> assignments = new HashMap<>();
        int parallelism = context.currentParallelism();

        for (MySplit split : pendingSplits) {
            int subtask = Math.abs(split.splitId().hashCode()) % parallelism;
            // 确保 subtask 是有效的并且已被注册
            if (context.registeredReaders().containsKey(subtask)) {
                assignments.computeIfAbsent(subtask, k -> new ArrayList<>()).add(split);
            } else {
                // 如果未注册，则保留这些分片等待下次分配
                System.err.println("Subtask " + subtask + " 尚未注册，暂时不分配分片");
            }
        }

        if (!assignments.isEmpty()) {
            context.assignSplits(new SplitsAssignment<>(assignments));
        }
        pendingSplits.clear();
    }


    /***
     * 处理请求分片的请求。这个方法的调用时机是在SourceReader通过提供subtaskId调用SourceReaderContext.sendSplitRequest()方法时。
     * 参数subtaskId，发送了source event的SourceReader的subtaskId。
     * 参数requesterHostname，非必填参数，发起分片请求的SourceReader任务的主机名，这个可以用来做分片下发时候的SourceReader的通信定位。
     * @param subtaskId
     * @param hostname
     */
    @Override
    public void handleSplitRequest(int subtaskId, @Nullable String hostname) {
        System.out.println("可在此实现按需分配（如根据 Reader 负载） subtaskId:"+subtaskId);
        assignSplits();
    }

    /**
     * 添加了一个新的SourceReader的subtaskId。
     * @param subtaskId
     */
    @Override
    public void addReader(int subtaskId) {
        System.out.println("新 Reader 注册时触发分配 subtaskId:"+subtaskId);
        assignSplits();
    }


    /**
     * 把下发给SourceReader的分片重新添加回枚举器，这个方法调用仅会发生在SourceReader Failover时，并且分配的分片列表是最近一次checkpoint成功的已分配分片列表。
     * 参数splits，重新给回枚举器待重新分配的分片列表。
     * 参数subtaskId，返回的这批分片原来归属的SourceReader task的subtaskId。
     * @param splits
     * @param subtaskId
     */
    @Override
    public void addSplitsBack(List<MySplit> splits, int subtaskId) {
        pendingSplits.addAll(splits); // 处理失败分片回收
    }


    /**
     * 生成枚举器的状态快照，并把状态快照当作结果返回，让其和checkpoint一起存储。
     * 快照应包含枚举器的最新状态：它应假设快照之前发生的所有操作都已成功完成。比如，分片已经分配给SourceReader(已经调用 SplitEnumeratorContext. assignSplit(SourceSplit, int) 和 SplitEnumeratorContext. assignSplits(SplitsAssignment) ) 就不再需要被包含在快照状态里了。
     * 参数checkpointId，状态快照的checkpoint id，大部分实现不关心这个参数，因为都比较关心快照的内容，而不太关心是哪次checkpoint产生的。这个参数对于有那些有外部系统的连接器会比较有用，因为连接器的外部系统会比较关心checkpoints。比如，枚举器可以通知连接器的外部系统一次特别的checkpoint已经被触发。
     * 返回参数，是一个包含了枚举器状态的对象。
     * throws Exception，当快照无法正确执行时。
     * @param checkpointId
     * @return
     */
    @Override
    public CheckpointState snapshotState(long checkpointId) {
        return new CheckpointState(pendingSplits); // 检查点状态快照
    }

    /**
     * 关闭枚举器，目的用来管理资源，比如线程、网络链接等。
     * 这个并不是枚举器原生接口，是其继承了AutoCloseable的接口。
     * @throws IOException
     */

    @Override
    public void close() throws IOException { }

    // 检查点状态序列化（需自定义）
    public static class CheckpointState {
        private final List<MySplit> remainingSplits;

        public CheckpointState(List<MySplit> remainingSplits) {
            this.remainingSplits = remainingSplits;
        }

        public Collection<? extends MySplit> getRemainingSplits() {
            return new ArrayList<>(remainingSplits.size());
        }
        // 构造方法及序列化逻辑...
    }


    /**
     * 接口功能描述：一个简单的版本序列化程序接口，序列化程序有一个版本（通过getVersion（）返回），可以附加到序列化数据上。当序列化程序发展时，版本可用于标识数据是使用哪个先前版本序列化的。
     * 泛型E，序列化程序进行序列化或者泛序列化的数据类型。
     */
    public static class CheckpointSerializer implements SimpleVersionedSerializer<CheckpointState> {
        @Override
        public int getVersion() {
            return 0;
        }

        @Override
        public byte[] serialize(CheckpointState checkpointState) throws IOException {
            return new byte[0];
        }

        @Override
        public CheckpointState deserialize(int i, byte[] bytes) throws IOException {
            return null;
        }
        // 实现序列化细节...
    }
}